import os
from zhipuai import ZhipuAI

from config import AI_KEY_ZHI_PU
from ai.common import process_messages, process_content

from ai.consts import PLATFORM_ZHI_PU


def get_result_by_zhipu(data, ai_info):
    client = ZhipuAI(api_key=AI_KEY_ZHI_PU)
    completion = client.chat.completions.create(
        model=ai_info['model'],
        messages=process_messages(data)
    )
    content = completion.choices[0].message.content
    return process_content(content, ai_info)
    
    
def get_zhipu_result(data, model_name):
    huoshan_ai_data = {
        "api_key": AI_KEY_ZHI_PU,
        "model": model_name,
        "platform": PLATFORM_ZHI_PU
    }
    return get_result_by_zhipu(data, huoshan_ai_data)

